from functools import cached_property
import logging
from pathlib import Path
import polars as pl
from pydantic import BaseModel
from pysam import VariantFile, VariantRecord  # pylint: disable=E0611

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("MTDD")


class Annotation(BaseModel):
    in_vcf_file: Path
    out_vcf_file: Path
    genepred_file: Path
    gene_file: Path

    @cached_property
    def transcripts(self) -> pl.DataFrame:
        t_records = (
            pl.read_csv(
                self.genepred_file,
                has_header=False,
                separator="\t",
                columns=[1, 2, 3, 4, 5, 6, 7, 9, 10, 12],
                new_columns=["transcript", "chrom", "strand", "tx_start", "tx_end", "cds_start", "cds_end", "exon_starts", "exon_ends", "gene"],
            )
            .with_columns(pl.col("exon_starts").str.strip_chars_end(",").str.split(",").cast(pl.List(pl.Int64)))
            .with_columns(pl.col("exon_ends").str.strip_chars_end(",").str.split(",").cast(pl.List(pl.Int64)))
            .with_columns(pl.when(pl.col("cds_end").eq(pl.col("cds_start"))).then(pl.lit("N")).otherwise(pl.lit("Y")).alias("type"))
            .select(["transcript", "chrom", "strand", "tx_start", "tx_end", "exon_starts", "exon_ends", "type", "gene"])
        )
        g_records = pl.read_csv(
            self.gene_file, separator="\t", has_header=False, columns=[0, 1, 2], new_columns=["chrom", "gene", "gene_id"]
        ).with_columns(pl.col("gene_id").cast(pl.String))
        return t_records.join(g_records, on=["chrom", "gene"], how="left").fill_null(".")

    def annotate(self, cnv: VariantRecord):
        detail_texts = []
        transcripts = self.transcripts.filter(pl.col("chrom").eq(cnv.chrom) & pl.col("tx_start").lt(cnv.stop) & pl.col("tx_end").ge(cnv.start))
        for (gene, transcript, strand, typo), records in transcripts.group_by(["gene", "transcript", "strand", "type"]):
            records = records.explode(["exon_starts", "exon_ends"]).rename({"exon_starts": "start", "exon_ends": "end"}).sort(["start"])
            e_records = (
                records.sort(["start"], descending=strand == "-")
                .with_row_index(offset=1)
                .with_columns(pl.format("exon{}", pl.col("index")).alias("region"))
            )
            i_records = (
                records[:-1]
                .with_columns(records["end"][:-1].alias("start"), e_records["start"][1:].alias("end"))
                .sort(["start"], descending=strand == "-")
                .with_row_index(offset=1)
                .with_columns(pl.format("region{}", pl.col("index")).alias("region"))
            )
            t_records = (
                pl.concat([e_records, i_records])
                .filter(pl.col("start").lt(cnv.stop) & pl.col("end").ge(cnv.start))
                .sort(["start"], descending=strand == "-")
            )
            region = t_records["region"][0] if t_records.shape[0] == 1 else f"{t_records['region'][0]}_{t_records['region'][-1]}"
            detail_texts.append(f"{gene}:{transcript}:{region}:{typo}")
        cnv.info.update(
            {
                "GENE": transcripts["gene"].to_list() or ["."],
                "GENE_ID": transcripts["gene_id"].to_list() or ["."],
                "DETAIL": detail_texts or ["."],
            }
        )

    def run(self):
        logger.info("Do Annotation to %s", self.out_vcf_file)
        with VariantFile(self.in_vcf_file) as reader:
            reader.header.info.add(id="GENE", number=".", type="String", description="Gene Symbol")
            reader.header.info.add(id="GENE_ID", number=".", type="String", description="Gene EntrezID")
            reader.header.info.add(id="DETAIL", number=".", type="String", description="Gene detail, FORMAT=Gene:Transcript:Region:RNA_type")
            reader.header.add_line("##source=MTDD")
            with VariantFile(self.out_vcf_file, "w", header=reader.header) as writer:
                for record in reader:
                    self.annotate(cnv=record)
                    writer.write(record)
